from TTS.api import TTS as CoquiTTS_API
from TTS.utils.radam import RAdam
import sounddevice as sd
import torch
import collections

# 添加 RAdam、collections.defaultdict 和 dict 到白名单
torch.serialization.add_safe_globals([RAdam, collections.defaultdict, dict])

# 初始化中文模型
print("正在加载中文模型...")
tts_zh = CoquiTTS_API("tts_models/zh-CN/baker/tacotron2-DDC-GST", progress_bar=False, gpu=False)
print("中文模型加载成功:", tts_zh.model_name)

# 初始化英文模型
print("正在加载英文模型...")
tts_en = CoquiTTS_API("tts_models/en/ljspeech/vits", progress_bar=False, gpu=False)
print("英文模型加载成功:", tts_en.model_name)

# 英文 TTS
text_en = "It took me quite a long time to develop a voice, and now that I have it I'm not going to be silent"
wav_en = tts_en.tts(text_en)
tts_en.synthesizer.save_wav(wav_en, "output_en.wav")
sd.play(wav_en, samplerate=tts_en.synthesizer.tts_config.audio["sample_rate"])  # 修正采样率访问
sd.wait()

# 中文 TTS
text_zh = "锄禾日当午，汗滴禾下土。"
wav_zh = tts_zh.tts(text_zh)
tts_zh.synthesizer.save_wav(wav_zh, "output_zh.wav")
sd.play(wav_zh, samplerate=tts_zh.synthesizer.tts_config.audio["sample_rate"])  # 修正采样率访问
sd.wait()